#include <cstdio>
#include <cstdlib>
//#include <tchar.h>

#include "opencv2/imgproc/imgproc.hpp"
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/video/video.hpp"
#include <opencv2/gpu/gpu.hpp>

#include "LcBasic.h"

#include "VideoRead.h"

#include "FeatureComputer.h"

#include "Classifier.h"

#include "Bagging.h"

#include "ClassifierRank.h"

#include "LcGraphCut.h"

#include "Test.h"

#include "InterFace.h"

#include "MotionAnalysis.h"

#include "TimeSpaceMRF.h"

class CollectInfo
{
public:
	vector<int> n_nodes;
	vector<int> node_shift;
	vector<Mat> node_descs;
	vector<Mat> label_maps;
	vector<SuperpixelFlow> superpixel_flows;
	vector<Mat> adj_mats;
	Mat gt;
};


LcValidator predictMRF(	vector< CollectInfo * > & collect_infos,
						Mat & parameter,						
						FILE * output = NULL,
						string * img_save_prefix = NULL);

void getReport(		string & report_name,
					Mat & parameter,
					vector< CollectInfo * > & collect_infos,
					string * save_name_prefix = NULL);

class TimeSpaceMRFOpt
{
public:

	TimeSpaceMRFOpt();

	TimeSpaceMRFOpt(vector< CollectInfo * >* );

	void set( vector< CollectInfo * >* );

	vector< CollectInfo * >* collect_infos;

	int jac_down_rate;

	float calcScore();

	void calcJacobian( Mat & jac);

	void work();

	Mat my_parameter;

	double learning_rate;
	double cooling_rate;
	int iter_times;

	float getScore(Mat & parameter);

	float default_lambda;
	float default_nu;
	float default_heat;
	float default_eta;

	bool FIX_LAMBDA;
	bool FIX_NU;
	bool FIX_HEAT;
	bool FIX_ETA;

	void init();

	Mat expandParameter( Mat & parameter);

};

TimeSpaceMRFOpt::TimeSpaceMRFOpt()
{
	jac_down_rate = 10;
	learning_rate = 0.4;
	cooling_rate = 0;
	iter_times = 100;
}

void TimeSpaceMRFOpt::set( vector< CollectInfo * >* n_collect_infos)
{
	collect_infos = n_collect_infos;
	my_parameter = Mat::ones( 4,  1 , 5) * 0.04;
	my_parameter.at<float>(3,0) = 0.02f;
}

TimeSpaceMRFOpt::TimeSpaceMRFOpt(vector< CollectInfo * >* n_collect_infos)
{
	learning_rate = 0.1;
	cooling_rate = 0;
	iter_times = 100;
	set(n_collect_infos);
	my_parameter = Mat::ones( 4, 1 , 5) * 0.04;	
}


float TimeSpaceMRFOpt::getScore(Mat & parameter)
{
	Mat exp_parameter = expandParameter( parameter );

	return predictMRF( *collect_infos,parameter).getF1();
}

void TimeSpaceMRFOpt::init()
{
	int dim = 0;
	if(!FIX_LAMBDA) dim++;
	if(!FIX_NU) dim++;
	if(!FIX_HEAT) dim++;
	if(!FIX_ETA) dim++;

	cout << dim << endl;

	my_parameter = Mat::zeros( dim , 1, 5 );

	float * p_ans = my_parameter.ptr<float>();

	if( !FIX_HEAT ){ *p_ans = 0.05f; p_ans++; }	

	if( !FIX_LAMBDA){ *p_ans = 0.1f; p_ans++; }	

	if( !FIX_NU){ *p_ans = 0.1f; p_ans++; }	

	if( !FIX_ETA){  *p_ans = 0.05f; p_ans++; }		
}

Mat TimeSpaceMRFOpt::expandParameter( Mat & parameter)
{
	float * p_data = parameter.ptr<float>();
	
	Mat exp_parameter = Mat::zeros(4,1,5);

	float * p_ans = exp_parameter.ptr<float>();

	if( FIX_HEAT ) *p_ans = default_heat;
	else{ *p_ans = *p_data; p_data++; }
	p_ans++;

	if( FIX_LAMBDA) *p_ans = default_lambda;
	else{ *p_ans = *p_data; p_data++; }
	p_ans++;

	if( FIX_NU)  *p_ans = default_nu;
	else{ *p_ans = *p_data; p_data++; }
	p_ans++;

	if( FIX_ETA)  *p_ans = default_eta;
	else{ *p_ans = *p_data; p_data++; }
	p_ans++;				

	return exp_parameter;
}

void TimeSpaceMRFOpt::work()
{
	CvTermCriteria term_criteria_ = CvTermCriteria();

	init();

	int dim = my_parameter.rows;	

	CvLevMarq solver(dim,
                     1,
                     term_criteria_);

	Mat err, jac;
	
	//my_parameter.convertTo( my_parameter, CV_64F );

    //CvMat matParams = my_parameter;
    //cvCopy(&matParams, solver.param);	

	float F_max = getScore(my_parameter);//Energy

	Mat my_parameter_max;

	my_parameter.copyTo(my_parameter_max);

	float temperature = 0.1f;

	for(int i = 0 ; i < iter_times; i++)
	{		

		double step = learning_rate
			* (1 - cooling_rate * double(i+1) / double(iter_times));

		calcJacobian(jac);		

		{
			//norm jac as a direction

			Mat temp = abs(jac);

			float norm = 1.0f/(1e-5f+lc::reduceMat( temp, CV_REDUCE_SUM));

			jac = jac * norm;
		}

		{

			int _dim = jac.cols;
			vector<float> score; score.resize(_dim );
			for(int i = 0 ; i < _dim ;i++)
				score[i] = -abs( jac.at<float>( 0, i ) );

			vector<int> id;
			lc::argsort( score, id);

			for(int i = 0; i < _dim ; i++)
				if(rand()/double(RAND_MAX) < double(i)/double(_dim)* 0.7 + 0.3)
					jac.at<float>(0,id[i]) = float(lc::normrnd()) * temperature * 0.1f;
		}		

		my_parameter = my_parameter - jac* step;		

		my_parameter = max(0,my_parameter);

		if(0){
			Mat temp = abs(my_parameter);
			float norm = 1.0f/(1e-5f+lc::reduceMat( temp, CV_REDUCE_SUM));
			my_parameter = my_parameter * norm;
		}

		float F_new = getScore( my_parameter);

		if(F_new > F_max)
		{
			F_max = getScore(my_parameter);//Energy	

			my_parameter.copyTo(my_parameter_max);
		}
		else if( rand()/double(RAND_MAX) > exp( -(F_max - F_new)/temperature))
		{
			my_parameter_max.copyTo(my_parameter);
		}

		temperature*=0.97f;

		if(i%1==0) cout << i+1 << "/" << iter_times << " " << "F= " << getScore(my_parameter) <<  " while para = " << my_parameter << endl;;
	}	

	my_parameter_max.copyTo(my_parameter);
}

void TimeSpaceMRFOpt::calcJacobian( Mat & jac)
{	
	my_parameter.convertTo( my_parameter, CV_32F );

	cout << my_parameter << endl;

	int dim = my_parameter.rows;

	//float step = 2e-2f;

	jac = Mat::zeros( dim, 1, 5);	

	float score_init = 1.0f - getScore( my_parameter );

	for(int i = 0 ; i < dim ; i ++)
	{		

		for(float step = 1e-3f; step< 0.2f; step*=1.5f)
		{
			Mat pos, neg;
			my_parameter.copyTo(pos);
			my_parameter.copyTo(neg);

			pos.at<float>(i,0) += step;
			neg.at<float>(i,0) -= step;		

			float score_pos = 1.0f - getScore( pos );
			float score_neg = 1.0f - getScore( neg );		

			if( score_pos< score_init && score_neg < score_init )
			{
				jac.at<float>(i,0) = (score_pos -score_neg) / 2.0f / step;
				//break;
			}
			else if(score_pos < score_init)
			{
				jac.at<float>(i,0) = (score_pos -score_init) / step * 2.0f;
				break;
			}
			else if(score_neg < score_init)
			{				
				jac.at<float>(i,0) = (score_init -score_neg) / step * 2.0f;
				break;
			}
			else
				jac.at<float>(i,0) = (score_pos -score_neg) / 2.0f / step * 0.5f;
		}
	}
	//jac = jac.t();
}

int main(int argc, char * const argv[])
{

	LcInterface my_setting;

	//char *const setting_string[] ={ "","-b","-v","lb","-t","GOPR2384_cab","-s", "500", "-e", "700", "-r","50","-p","GOPR2382_cab","-f","./../../_FGS/","-w","2"};

	if(argc==1)
	{
		char *const setting_string[] ={ "","-b","-v","lb","-t","GOPR2384_cab","-s", "500", "-e", "510", "-r","20","-p","GOPR2393_cab","-f","./../../_FGS/","-w","5","-z","_test","-i","0.1101","-j","0.1086","-k","0.0725","-l","2","-O"};

		my_setting.input( 28 + 1 ,setting_string);
	}
	else my_setting.input( argc, argv );

	ClassifierRank my_crank;

	LcVideoReadExt my_reader = my_setting.getTrainVideo();

	my_crank.setReader( &my_reader);

	my_crank.setFeatureCode( my_setting.feature_code_string.c_str() );

	my_crank.setClassifier( my_setting.getClassifier() );

	vector< CollectInfo * > collect_infos;

	if(my_setting.IF_TEST)
	{
		LocalSensitiveBagging crank_bagging;

		crank_bagging.setFeatureCode(my_setting.feature_code_string.c_str());

		crank_bagging.setReader( &my_reader );

		{
			string saving_name = my_reader.root + "classifier_output/crank/";
			crank_bagging.load( saving_name.c_str() );
		}

		LcHistogram global_extractor;

		LcVideoReadExt test_video = my_setting.getTestVideo();		

		for(int i = 0; i < 1e4;i++)
		{
			Mat frm;
			vector<Mat> frm_array;
			int statu = test_video.readNext( frm , &frm_array);

			if(statu!=0) break;

			int frame_id = test_video.getf();

			cout << "=================" << frame_id << "===========" << endl;

			Mat gt;
			test_video.getLabel( gt );
			if(!gt.data){cout << "warning! gt unfound " << endl; continue;}
			resize( gt,gt,frm.size() );

			vector< Mat > ress;

			int filter_rate = my_setting.frame_filter_rate;

			int j = 0;
			for(int k = 0 ; k < frm_array.size();k++)
			{
				if( (k - my_setting.t_win) % filter_rate==0)
				{
					frm_array[j] = frm_array[k];
					j++;
				}
			}
			frm_array.resize( j );

			if(my_setting.SAVE_LIKELI)
			{
				int t_win = int( frm_array.size() ) /2;
				string save_name_prefix = my_setting.test_video_name;

				Mat res;
				Mat emptygt = Mat::Mat();

				crank_bagging.predictFrame( frm_array[ t_win], res, emptygt, global_extractor);

				if(res.type() == 5)
				{
					res = res * 255.0;
					res.convertTo( res, CV_8U );
				}

				char fname[400];
				sprintf( fname , "./save_img/%s_likeli_%d.jpg", save_name_prefix.c_str() , i);
				imwrite(fname , res);
				sprintf( fname , "./save_img/%s_src_%d.jpg", save_name_prefix.c_str() , i);
				imwrite(fname, frm_array[ t_win]);
				continue;
			}

			CollectInfo * frame_info = new CollectInfo;

			collect_infos.push_back( frame_info );

			{
				TimeSpaceMRF mrf_get;

				mrf_get.getAllInfo(	frm_array, crank_bagging, global_extractor,
									frame_info->n_nodes,
									frame_info->node_shift,
									frame_info->node_descs,
									frame_info->label_maps,
									frame_info->superpixel_flows,
									frame_info->adj_mats );

				gt.copyTo( frame_info->gt );
			}			

			


		}// end for video

		if(my_setting.SAVE_LIKELI) return 0;


		if(my_setting.SAVE_IMG)
		{
			{
				string report_name = "./report/Valid_" + my_setting.test_video_name + "_on_"
					+ my_setting.train_video_name + my_setting.addition_string + ".txt";			

				string save_name_prefix = my_setting.test_video_name + my_setting.addition_string;

				Mat parameter = Mat::zeros(4,1,5);

				parameter.at<float>(0,0) = my_setting.weight_heat_map;
				parameter.at<float>(0,1) = my_setting.lambda;
				parameter.at<float>(0,2) = my_setting.nu;
				parameter.at<float>(0,3) = my_setting.eta;

				getReport(report_name, parameter, collect_infos, & save_name_prefix);
			}
		}
		else
		{
			{
				string report_name = "./report/JointPredict_" + my_setting.test_video_name + "_on_"
					+ my_setting.train_video_name + my_setting.addition_string + ".txt";			

				Mat parameter;

				getReport(report_name, parameter, collect_infos);
			}

			TimeSpaceMRFOpt mrf_opt( &collect_infos);

			mrf_opt.learning_rate = my_setting.learning_rate;
			mrf_opt.iter_times = my_setting.iter_times;
			mrf_opt.cooling_rate = my_setting.cooling_rate;

			mrf_opt.FIX_ETA = my_setting.FIX_ETA;
			mrf_opt.FIX_HEAT = my_setting.FIX_HEAT;
			mrf_opt.FIX_LAMBDA = my_setting.FIX_LAMBDA;
			mrf_opt.FIX_NU = my_setting.FIX_NU;
			

			mrf_opt.work();
			
			Mat parameter_opt = mrf_opt.expandParameter( mrf_opt.my_parameter );

			{
				string report_name = "./report/Opt_" + my_setting.test_video_name + "_on_"
					+ my_setting.train_video_name + my_setting.addition_string + ".txt";						

				getReport(report_name, parameter_opt, collect_infos);

				string para_save_name = "./report/para_" + my_setting.test_video_name + "_on_"
					+ my_setting.train_video_name + my_setting.addition_string + ".bin";

				lc::LcMat2Bin( para_save_name.c_str() , parameter_opt);
			}
		}

	}	

	return 0;
}


LcValidator predictMRF(	vector< CollectInfo * > & collect_infos,
						Mat & parameter,						
						FILE * fout_report,
						string * img_save_prefix)
{
	LcValidator sum_score(0,0,0,0);

	for(int i = 0 ; i < (int) collect_infos.size(); i++)
	{
		CollectInfo * frame_info = collect_infos[i];

		{
			vector< Mat > ress;

			TimeSpaceMRF mrf_test;			

			if(parameter.rows==4)
			{
				mrf_test.heat_map_weight = parameter.at<float>(0,0);
				mrf_test.lambda          = parameter.at<float>(1,0);
				mrf_test.nu              = parameter.at<float>(2,0);
				mrf_test.eta			 = parameter.at<float>(3,0);

			}			

			LcValidator _score = mrf_test.predict(
								frame_info->n_nodes,
								frame_info->node_shift,
								frame_info->node_descs,
								frame_info->label_maps,
								frame_info->superpixel_flows,
								frame_info->adj_mats,
								frame_info->gt, ress);

			if(img_save_prefix!=NULL)
			{
				int t_win = int(ress.size()) /2;
				char fname[400];
				sprintf( fname , "./save_img/%s_%d.jpg", img_save_prefix->c_str() , i);

				Mat dsp;

				if(ress[t_win].type() == 5)
				{
					dsp = ress[t_win]*255.0;
					dsp.convertTo(dsp,CV_8U);
				}
				else dsp = ress[t_win];

				imwrite( fname , dsp );
			}

			//_score.display();

			if(fout_report!=NULL)
				fprintf(fout_report, "%f\t%f\t%f\t%f\t%f\n", _score.getF1(),
					_score.tp,_score.fp,_score.fn,_score.tn);

			sum_score = sum_score + _score;
			
		}
	}

	return sum_score;
}

void getReport( string & report_name, Mat & parameter, vector< CollectInfo * > & collect_infos, string * save_name_prefix)
{
	FILE * fout_report;

	fout_report = fopen(report_name.c_str(),"w");

	LcValidator sum_score = predictMRF(	collect_infos,
										parameter,											
										fout_report,
										save_name_prefix);

	fprintf(fout_report, "-1\t-1\t-1\t-1\t-1\n");

	fprintf(fout_report, "%f\t%f\t%f\t%f\t%f\n", sum_score.getF1(),
				sum_score.tp,sum_score.fp,sum_score.fn,sum_score.tn);

	fclose(fout_report);
}